import base64
import json
import os
import requests

folderPath = 'C:\\Users\\sgfirstliu\\Documents\WXWork\\1688856156362307\Cache\File\\2024-05\SM_Roman_NPC_Level01\\'  # 把目录改成自己的

allFilename = []
def file_to_base64(filename):
    allFilename.append(filename)
    with open(folderPath + filename, "rb") as file:
        data = file.read()
        base64_str = str(base64.b64encode(data), "utf-8")
        return "data:image/png;base64," + base64_str

def get_img2img_param(img2img_basic_image_base64, denoising_strength, width, height):
    return {
        "batch_size": 1,
        "cfg_scale": 7,
        "denoising_strength": denoising_strength,
        "eta": 0,
        "include_init_images": False,
        "init_images": [img2img_basic_image_base64],
        "inpaint_full_res": False,
        "inpaint_full_res_padding": 0,
        "inpainting_fill": 0,
        "inpainting_mask_invert": False,
        "mask": None,
        "mask_blur": 4,
        "n_iter": 1,
        "negative_prompt": "",
        "override_settings": {},
        "prompt": "example prompt",
        "resize_mode": 0,
        "restore_faces": False,
        "s_churn": 0,
        "s_noise": 1,
        "s_tmax": 0,
        "s_tmin": 0,
        "sampler_index": "Euler a",
        "seed": -1,
        "seed_resize_from_h": -1,
        "seed_resize_from_w": -1,
        "steps": 3,
        "styles": [],
        "subseed": -1,
        "subseed_strength": 0,
        "tiling": False,
        "width": width,
        "height": height,
    }

def url_img2img(base_url):
    return f"{base_url}/sdapi/v1/img2img"

# 把URL都改成自己的，参数也都可以调
img2img_basic_image_base64 = file_to_base64('T_RomanPalazzo01_D.PNG')
rsp = requests.post(url_img2img('http://127.0.0.1:7860'), json=get_img2img_param(img2img_basic_image_base64, 0.3, 1024, 1024))
if rsp.status_code != 200:
    print('connect failed')
else:
    retObj = json.loads(rsp.text)
    print(retObj.keys())
    for index, imgBase64 in enumerate(retObj['images']): # 保存所有图片
        imgBinary = base64.b64decode(imgBase64)
        print(folderPath + 'new/' + allFilename[index])
        with open(folderPath + 'new/' + allFilename[index], "wb") as file:
            file.write(imgBinary)